import argparse
import logging
from dingtalk_stream import AckMessage
import dingtalk_stream
from dotenv import load_dotenv
import os
import asyncio
import settings
from aiagent.agent import ZLAgent


load_dotenv()


# clent = ZLAgent()
# await clent.initialize()

def setup_logger():
    logger = logging.getLogger()
    handler = logging.StreamHandler()
    handler.setFormatter(
        logging.Formatter('%(asctime)s %(name)-8s %(levelname)-8s %(message)s [%(filename)s:%(lineno)d]'))
    logger.addHandler(handler)
    logger.setLevel(logging.INFO)
    return logger


# class EchoTextHandler(dingtalk_stream.ChatbotHandler):
#     def __init__(self, logger: logging.Logger = None):
#         super(dingtalk_stream.ChatbotHandler, self).__init__()
#         if logger:
#             self.logger = logger
#
#     async def process(self, callback: dingtalk_stream.CallbackMessage):
#         print("callback:::", callback)
#         print("callback.data:::", callback.data)
#         senderStaffId = callback.data.get("senderStaffId")
#         print(f"senderStaffId:::, {senderStaffId}")
#         incoming_message = dingtalk_stream.ChatbotMessage.from_dict(callback.data)
#         text = incoming_message.text.content.strip()
#         self.reply_text(text, incoming_message)
#         return AckMessage.STATUS_OK, 'ok'


class EchoTextHandler(dingtalk_stream.ChatbotHandler):
    def __init__(self, logger: logging.Logger = None):
        super(dingtalk_stream.ChatbotHandler, self).__init__()
        if logger:
            self.logger = logger
        self.agent = None  # 延迟初始化

    # async def process(self, callback: dingtalk_stream.CallbackMessage):
    #     incoming_message = dingtalk_stream.ChatbotMessage.from_dict(callback.data)
    #     text = (incoming_message.text.content or "").strip()
    #     sender_staff_id = callback.data.get("senderStaffId")
    #
    #     if not text:
    #         self.reply_text("抱歉，我没听清楚你说什么", incoming_message)
    #         return AckMessage.STATUS_OK, 'OK'
    #
    #     if self.agent is None or self.agent.agent_executor is None:
    #         try:
    #             self.agent = ZLAgent()
    #             await self.agent.initialize()
    #             self.logger.info("✅ ZLAgent 初始化成功")
    #         except Exception as e:
    #             self.logger.error(f"❌ 初始化失败: {e}", exc_info=True)
    #             self.reply_text("系统错误，请稍后重试。", incoming_message)
    #             return AckMessage.STATUS_OK, 'OK'
    #
    #     # --- 流式发送逻辑 ---
    #     buffer = ""  # 累积 token
    #     final_output = ""  # 完整回复
    #
    #     try:
    #         # 使用 astream_events 获取事件流
    #         async for event in self.agent.agent_executor.astream_events(
    #                 {
    #                     "input": text,
    #                     "sender_staff_id": sender_staff_id
    #                 },
    #                 config={"configurable": {"session_id": sender_staff_id}},
    #                 version="v2"  # 必须是 v2
    #         ):
    #             # 只关注 LLM 生成的 token
    #             if event["event"] == "on_chat_model_stream":
    #                 content = event["data"].get("chunk", {}).content
    #                 if content:
    #                     buffer += content
    #                     final_output += content
    #
    #                     # 判断是否该发送
    #                     if content in '。！？\n' and len(buffer.strip()) > 10:
    #                         self.reply_text(buffer.strip(), incoming_message)
    #                         buffer = ""
    #                         await asyncio.sleep(0.5)  # 模拟打字速度
    #
    #         # 发送最后剩下的内容
    #         if buffer.strip():
    #             self.reply_text(buffer.strip(), incoming_message)
    #             buffer = ""
    #
    #         # 把完整回复存入 history
    #         from aiagent.agent import get_session_history
    #         history = get_session_history(sender_staff_id)
    #         history.add_ai_message(final_output)
    #
    #     except Exception as e:
    #         self.logger.error(f"❌ 流式运行失败: {e}", exc_info=True)
    #         self.reply_text("处理失败，请稍后再试。", incoming_message)
    #
    #     return AckMessage.STATUS_OK, 'OK'


    async def process(self, callback: dingtalk_stream.CallbackMessage):
        print("callback:::", callback)
        print("callback.data:::", callback.data)
        sender_staff_id = callback.data.get("senderStaffId")
        print(f"senderStaffId::: {sender_staff_id}")

        incoming_message = dingtalk_stream.ChatbotMessage.from_dict(callback.data)
        text = (incoming_message.text.content or "").strip()

        if not text:
            self.reply_text("抱歉，我没听清楚你说什么", incoming_message)
            return AckMessage.STATUS_OK, 'OK'

        # --- 关键：延迟初始化 ZLAgent ---
        if self.agent is None:
            try:
                self.agent = ZLAgent()
                await self.agent.initialize()
                self.logger.info("✅ ZLAgent 初始化成功")
            except Exception as e:
                self.logger.error(f"❌ ZLAgent 初始化失败: {e}")
                self.reply_text("系统初始化失败，请稍后重试。", incoming_message)
                return AckMessage.STATUS_OK, 'OK'

        # --- 运行 Agent ---
        try:
            response = await self.agent.run(text, sender_staff_id=sender_staff_id)
            if isinstance(response, dict):
                reply_content = response.get('output', str(response))
            else:
                reply_content = str(response)
            self.reply_text(reply_content, incoming_message)
        except Exception as e:
            self.logger.error(f"❌ Agent 运行失败: {e}", exc_info=True)
            self.reply_text("抱歉，我出了一点问题，请稍后再试。", incoming_message)

        return AckMessage.STATUS_OK, 'OK'




def main():
    logger = setup_logger()
    credential = dingtalk_stream.Credential(os.getenv('DING_Client_ID'), os.getenv('DING_Client_Secret'))
    client = dingtalk_stream.DingTalkStreamClient(credential)
    client.register_callback_handler(dingtalk_stream.chatbot.ChatbotMessage.TOPIC, EchoTextHandler(logger))
    client.start_forever()

if __name__ == '__main__':
    asyncio.run(main())


# async def gaodeceshi():
#     clent = ZLAgent()
#     await clent.initialize()
#     while True:
#         result = await clent.run("现在美国时间是几点", sender_staff_id="2611080337-1837672988")
#         print(result)
#         await asyncio.sleep(10)
#
#
# if __name__ == '__main__':
#     asyncio.run(gaodeceshi())
